import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import random
import jieba
from collections import Counter
from sklearn.metrics.pairwise import cosine_similarity

# ========== 1. 数据预处理 ==========
corpus = [
    "我们 喜欢 深度 学习",
    "自然 语言 处理 是 有趣 的",
    "人工智能 改变 了 世界",
    "深度 学习 是 人工智能 的 重要 组成部分"
]

# 超参数
window_size = 2  # 窗口大小
embedding_dim = 10  # 词向量维度
num_epochs = 100  # 训练轮数
learning_rate = 0.01  # 学习率
batch_size = 4  # 批大小
neg_samples = 5  # 负采样个数

# 分词 & 构建词汇表
tokenized_corpus = [list(jieba.cut(sentence)) for sentence in corpus]
vocab = set(word for sentence in tokenized_corpus for word in sentence)
word2idx = {word: idx for idx, word in enumerate(vocab)}
idx2word = {idx: word for word, idx in word2idx.items()}

# 统计词频
word_counts = Counter([word for sentence in tokenized_corpus for word in sentence])
total_words = sum(word_counts.values())

# 计算负采样概率
word_freqs = {word: count / total_words for word, count in word_counts.items()}
word_powers = {word: freq ** 0.75 for word, freq in word_freqs.items()}
Z = sum(word_powers.values())
word_distribution = {word: prob / Z for word, prob in word_powers.items()}


# 负采样函数
def negative_sampling(positive_word, num_samples=5):
    words = list(word_distribution.keys())
    probabilities = list(word_distribution.values())
    negatives = []
    while len(negatives) < num_samples:
        neg = np.random.choice(words, p=probabilities)
        if neg != positive_word:
            negatives.append(neg)
    return negatives


# 生成 Skip-gram 训练数据
data = []
for sentence in tokenized_corpus:
    indices = [word2idx[word] for word in sentence]
    for center_idx in range(len(indices)):
        center_word = indices[center_idx]
        for offset in range(-window_size, window_size + 1):
            context_idx = center_idx + offset
            if 0 <= context_idx < len(indices) and context_idx != center_idx:
                context_word = indices[context_idx]
                data.append((center_word, context_word))

# 转换为 PyTorch 张量
data = [(torch.tensor(center), torch.tensor(context)) for center, context in data]


# ========== 2. 定义 Word2Vec (Skip-gram) 模型 ==========
class Word2Vec(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(Word2Vec, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.output_layer = nn.Linear(embedding_dim, vocab_size)

    def forward(self, center_word):
        embed = self.embedding(center_word)  # 获取中心词向量
        out = self.output_layer(embed)  # 计算词分布
        return out


# 初始化模型
model = Word2Vec(len(vocab), embedding_dim)

# ========== 3. 训练 Word2Vec ==========
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    total_loss = 0
    random.shuffle(data)  # 每轮打乱数据
    for center_word, context_word in data:
        optimizer.zero_grad()
        output = model(center_word.unsqueeze(0))  # 预测词分布
        loss = criterion(output, context_word.unsqueeze(0))  # 计算损失
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    if (epoch + 1) % 10 == 0:
        print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {total_loss:.4f}")

# ========== 4. 测试词向量 ==========
word_vectors = model.embedding.weight.data.numpy()


# 计算单词相似度
def most_similar(word, top_n=3):
    if word not in word2idx:
        return "单词不在词汇表中"

    word_vec = word_vectors[word2idx[word]].reshape(1, -1)
    similarities = cosine_similarity(word_vec, word_vectors)[0]

    # 获取相似度最高的 top_n 个单词（排除自身）
    similar_idx = similarities.argsort()[::-1][1:top_n + 1]
    return [(idx2word[idx], similarities[idx]) for idx in similar_idx]


# 测试相似词
test_words = ["深度", "学习", "人工智能"]
for word in test_words:
    print(f"【{word}】的相似单词:", most_similar(word))
